import pandas as pd
from pathlib import Path

def analyze_transaction_types():
    """分析银行流水数据，根据借贷标志和交易摘要进行分类"""
    try:
        # 配置数据文件路径
        data_dir = Path('data')
        source_dir = data_dir / 'source'
        
        # 遍历source目录下的所有子文件夹
        if source_dir.exists() and source_dir.is_dir():
            for folder in source_dir.iterdir():
                if folder.is_dir():
                    process_folder(folder)
        else:
            print("未找到source目录")
            return False
        
        return True
    except Exception as e:
        print(f"分析过程中出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

def process_folder(folder_path):
    """处理单个文件夹的数据"""
    try:
        print(f"\n处理文件夹: {folder_path.name}")
        
        # 读取银行流水数据
        bank_file = folder_path / "银行流水数据库.xlsx"
        if not bank_file.exists():
            print(f"银行流水文件不存在: {bank_file}")
            return
        
        # 读取亲属关系数据
        relative_file = folder_path / "亲属关系.xlsx"
        if not relative_file.exists():
            print(f"亲属关系文件不存在: {relative_file}")
            return
        
        # 读取Excel文件
        bank_df = pd.read_excel(bank_file)
        relative_df = pd.read_excel(relative_file, header=1)  # 从第二行开始读取
        
        # 获取亲属关系表中的名称列表
        relative_names = set(relative_df['名称'].dropna().unique())
        
        # 筛选银行流水中名称与亲属关系表匹配的数据
        filtered_df = bank_df[bank_df['名称'].isin(relative_names)].copy()
        
        if filtered_df.empty:
            print(f"未找到匹配的交易记录")
            return
        
        # 添加类型列
        filtered_df['类型'] = filtered_df['借贷标志'].apply(categorize_transaction_type)
        
        # 添加详细类型列
        filtered_df['详细类型'] = filtered_df.apply(
            lambda row: categorize_detailed_type(row['交易摘要'], row['类型']), 
            axis=1
        )
        
        # 创建输出目录
        output_dir = Path('data') / 'output' / folder_path.name
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # 保存结果
        output_file = output_dir / "交易分类结果.xlsx"
        filtered_df.to_excel(output_file, index=False)
        print(f"分析结果已保存到: {output_file}")
        print(f"共处理 {len(filtered_df)} 条记录")
        
    except Exception as e:
        print(f"处理文件夹 {folder_path.name} 时出错: {str(e)}")

def categorize_transaction_type(flag):
    """根据借贷标志分类交易类型"""
    if flag == '进':
        return '收入'
    elif flag == '出':
        return '一般支出'
    else:
        return '其他'

def categorize_detailed_type(summary, type_):
    """根据交易摘要和类型分类详细类型"""
    if not isinstance(summary, str):
        return '其他'
        
    summary = summary.lower()
    
    if type_ == '收入':
        if '理财' in summary:
            return '理财'
        elif '股票' in summary:
            return '股票'
        elif '工资' in summary:
            return '工资'
        else:
            return '其他'
    elif type_ == '一般支出':
        if '支付宝' in summary:
            return '支付宝'
        elif '取现' in summary:
            return '取现'
        elif '购物' in summary:
            return '购物'
        else:
            return '其他'
    else:  # 类型为其他
        return '其他'

if __name__ == "__main__":
    print("开始分析交易类型...")
    success = analyze_transaction_types()
    if success:
        print("分析完成！")
    else:
        print("分析失败！") 